/-
Copyright (c) 2025 Lean FRO, LLC. All rights reserved.
Released under Apache 2.0 license as described in the file LICENSE.
Authors: Joseph Rotella
-/

module

prelude
public import Lean.Meta.Hint

public section

/-!
# Missing-field hints for structure instance elaborator

This module generates suggested edits for structure instance notation that is missing required
fields. The suggestions adapt to the style of the existing structure instance syntax and
automatically insert any values inferred via constraints generated by already present fields.
-/

namespace Lean.Elab.Term.StructInst

open Meta
open TSyntax.Compat

/-- Auxiliary type for `mkMissingFieldsHint` -/
private structure FieldsHintView where
  /-- The position of the first existing field in the structure instance, if one exists. -/
  initFieldPos? : Option String.Pos.Raw
  /-- The tail position of the last existing field, or `none` if there are none. -/
  lastFieldTailPos? : Option String.Pos.Raw
  /-- Does this structure instance have a `with` clause? -/
  hasWith : Bool
  /-- The number of existing fields present in the structure instance. -/
  numFields : Nat
  /-- The position of the opening syntax (i.e., `{`) for the structure instance. -/
  openingPos : String.Pos.Raw
  /-- The position of the "leader" immediately preceding the fields -- either `with` or `{`. -/
  leaderPos : String.Pos.Raw
  /-- The tail position of the "leader" syntax. -/
  leaderTailPos : String.Pos.Raw
  /--
  The position of the closing syntax (i.e., `}`) for the structure instance.

  We refer only to the "closing" token `}` rather than a flexible "trailer" because macro-expansion
  erases the type annotation if it was present, and it's impossible to have an ellipsis if we're
  generating a "fields missing" error.
  -/
  closingPos : String.Pos.Raw

/--
Creates a hint message for the "fields missing" error that fills in the missing fields, adapting to
the structure instance notation style of the current syntax. Note that if the syntax is malformed or
synthetic, this function returns an empty message instead.
-/
def mkMissingFieldsHint (fields : Array (Name × Option Expr)) (stx : Syntax) : MetaM MessageData := do
  let some view := mkFieldsHintView? stx | return .nil
  let newFields ← fields.mapM fun (fieldName, fieldVal?) => do
    let value ← if let some fieldVal := fieldVal? then
      let stx ← withOptions (pp.mvars.set · false) <| PrettyPrinter.delab fieldVal
      pure (← PrettyPrinter.ppCategory `term stx).pretty
    else pure "_"
    pure s!"{fieldName} := {value}"
  let useSingleLine ← isSingleLineStyle view
  let suggestionText :=
    if useSingleLine then
      .group (behavior := .fill) <|
        newFields.foldl (init := Std.Format.nil) (fun acc sug => acc ++ "," ++ .line ++ format sug)
    else
      Format.line.joinSep newFields.toList

  let fileMap ← getFileMap
  let col (p : String.Pos.Raw) := (fileMap.utf8PosToLspPos p).character
  let indent :=
    if let some initFieldPos := view.initFieldPos? then
      col initFieldPos
    else
      -- Take the minimum of the opening and closing brace indentations to account for both the
      -- `...\n{ f := v\n ... }` and `...{\n  f := v\n ... \n}` styles
      min (col view.openingPos) (col view.closingPos) + 2

  let leaderLine := (fileMap.utf8PosToLspPos view.leaderPos).line
  let leaderClosingSeparated : Bool := leaderLine < (fileMap.utf8PosToLspPos view.closingPos).line
  -- If a line separates the leader and `}`, and there's no intervening text, `interveningLineEndPos?`
  -- gives the position on that line (the min of column `indent` and the line end) at which to insert
  let interveningLineEndPos? :=
    if leaderLine + 1 ≥ (fileMap.utf8PosToLspPos view.closingPos).line then none else
      let leaderLineEnd := findLineEnd fileMap.source view.leaderTailPos
      let indentPos := (indent + 1).fold (init := leaderLineEnd) (fun _ _ p => p.next fileMap.source)
      let interveningLineEnd := findLineEnd fileMap.source (leaderLineEnd.next fileMap.source)
      let nextTwoLines := Substring.Raw.mk fileMap.source view.leaderTailPos interveningLineEnd
      if nextTwoLines.all (·.isWhitespace) then some (min indentPos nextTwoLines.stopPos) else none

  let (preWs, postWs) : Format × Format :=
    -- If we're in single-line mode, `suggestionText` already has a leading `, `, so insert it as-is
    if useSingleLine then
      (.nil, .nil)
    -- Ensure a line break for consistent indentation if fields exist or there are no fields after `with`
    else if view.numFields > 0 || (view.hasWith && !leaderClosingSeparated) then
      (.line, .nil)
    -- If there is no intervening line, add spaces before and after for `{ f := v\n ...}` style
    else if !leaderClosingSeparated then
      (" ", " ")
    -- If there is an intervening line, insert therein if the region is empty; else insert before
    -- it, as the nonempty region may contain a type annotation
    else
      match interveningLineEndPos? with
      | none => (.line, .nil)
      | some interveningLineEndPos =>
        (String.ofList (List.replicate (indent - col interveningLineEndPos) ' '), .nil)
  let suggestionText := preWs ++ suggestionText ++ postWs
  let insPos := view.lastFieldTailPos?.getD <| interveningLineEndPos?.getD view.leaderTailPos
  let width := Tactic.TryThis.format.inputWidth.get (← getOptions)
  -- Limit only the span of the suggestion so the code action appears everywhere in the structure
  let suggestion := {
    suggestion := suggestionText.pretty width indent (col insPos),
    span? := Syntax.ofRange ⟨insPos, insPos⟩,
    toCodeActionTitle? := some fun _ => "Add missing fields"
  }
  MessageData.hint m!"Add missing fields:" #[suggestion]
where
  /--
  Constructs a `FieldsHintView` used to generate missing-fields hint suggestions. This method
  consolidates the low-level syntax inspection required for `mkMissingFieldsHint`.
  -/
  mkFieldsHintView? (stx : Syntax) : Option FieldsHintView := do
    -- Only show suggestions for user-entered structure-instance notation
    -- TODO: selectively let through macro-expanded `where` notation
    guard <| stx.getHeadInfo matches .original ..
    guard <| stx.getKind == ``Parser.Term.structInst
    let (leader, hasWith) := if stx[1][1] matches .missing then
      (stx[0], false)
    else
      (stx[1][1], true)
    let fields := stx[2][0].getSepArgs
    some { leaderPos := (← leader.getPos?)
           leaderTailPos := (← leader.getTailPos?)
           openingPos := (← stx[0].getPos?)
           closingPos := (← stx[5].getPos?)
           numFields := fields.size
           initFieldPos? := fields[0]?.bind (·.getPos?)
           lastFieldTailPos? := fields[fields.size - 1]?.bind (·.getTailPos?)
           hasWith }

  /--
  Returns the position of the end of the line that contains the position `pos`.
  Counterpart of `String.findLineStart`.
  -/
  findLineEnd (s : String) (p : String.Pos.Raw) : String.Pos.Raw :=
    (s.pos! p).find '\n' |>.offset

  /--
  Is the structure instance notation `stx` described by `view` using single-line styling?

  We infer this based on whether the last two fields are comma-separated and either on the same line
  or indented in a way not possible with one-field-per-line styling.
  -/
  isSingleLineStyle (view : FieldsHintView) : MetaM Bool := do
    unless view.numFields ≥ 2 do return false
    let rawFields := stx[2][0]
    -- Account for the fact that we may have a trailing separator:
    let lastInterveningSepIdx := rawFields.getNumArgs - 2 - ((rawFields.getNumArgs + 1) % 2)
    unless rawFields[lastInterveningSepIdx].getKind == `«,» do
      return false

    let some penultimateFieldPos := rawFields[lastInterveningSepIdx - 1].getPos? | return false
    let some lastFieldPos := rawFields[lastInterveningSepIdx + 1].getPos? | return false
    let some lastFieldTailPos := rawFields[lastInterveningSepIdx + 1].getTailPos? | return false
    let fileMap ← getFileMap
    let ⟨penultimateFieldLine, penultimateFieldCol⟩ := fileMap.utf8PosToLspPos penultimateFieldPos
    let ⟨lastFieldLine, lastFieldCol⟩ := fileMap.utf8PosToLspPos lastFieldPos
    return lastFieldLine == penultimateFieldLine || lastFieldCol < penultimateFieldCol
